// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/sparksql/specialforms/make_decimal.h>
#include <pollux/expression/constant_expr.h>

namespace kumo::pollux::functions::sparksql {
    namespace {
        template<typename T>
        class MakeDecimalFunction final : public exec::VectorFunction {
        public:
            MakeDecimalFunction(uint8_t resultPrecision, bool nullOnOverflow)
                : resultPrecision_(resultPrecision), nullOnOverflow_(nullOnOverflow) {
            }

            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args, // Not using const ref so we can reuse args
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const final {
                POLLUX_USER_CHECK_GE(
                    args.size(), 1, "Make_decimal expects at least one argument.");
                POLLUX_USER_CHECK(
                    args[0]->is_constant_encoding() || args[0]->is_flat_encoding(),
                    "Single-arg deterministic functions receive their only argument as flat or constant vector.");
                context.ensure_writable(rows, outputType, result);
                result->clear_nulls(rows);
                auto rawResults = result->as_unchecked<FlatVector<T> >()->mutableRawValues();
                const int128_t bound = DecimalUtil::kPowersOfTen[resultPrecision_];
                if (args[0]->is_constant_encoding()) {
                    // Fast path for constant vectors.
                    applyConstant(rows, args[0], bound, context, result, rawResults);
                } else {
                    // Fast path for flat vectors.
                    applyFlat(rows, args[0], bound, context, result, rawResults);
                }
            }

        private:
            void applyConstant(
                const SelectivityVector &rows,
                VectorPtr &arg,
                int128_t bound,
                exec::EvalCtx &context,
                VectorPtr &result,
                T *rawResults) const {
                const auto constant =
                        arg->as_unchecked<ConstantVector<int64_t> >()->value_at(0);
                if constexpr (std::is_same_v<T, int64_t>) {
                    const bool overflow = constant <= -bound || constant >= bound;
                    if (overflow) {
                        // The unscaled value cannot fit in the requested precision.
                        if (nullOnOverflow_) {
                            result->add_nulls(rows);
                        } else {
                            context.setErrors(
                                rows, overflowException(constant, resultPrecision_));
                        }
                    } else {
                        rows.applyToSelected(
                            [&](vector_size_t row) { rawResults[row] = constant; });
                    }
                } else {
                    rows.applyToSelected([&](vector_size_t row) {
                        rawResults[row] = static_cast<int128_t>(constant);
                    });
                }
            }

            void applyFlat(
                const SelectivityVector &rows,
                VectorPtr &arg,
                int128_t bound,
                exec::EvalCtx &context,
                VectorPtr &result,
                T *rawResults) const {
                auto rawValues =
                        arg->as_unchecked<FlatVector<int64_t> >()->mutableRawValues();
                if constexpr (std::is_same_v<T, int64_t>) {
                    context.applyToSelectedNoThrow(rows, [&](auto row) {
                        const auto unscaled = rawValues[row];
                        if (unscaled <= -bound || unscaled >= bound) {
                            // The unscaled value cannot fit in the requested precision.
                            if (nullOnOverflow_) {
                                result->set_null(row, true);
                            } else {
                                context.setError(
                                    row, overflowException(unscaled, resultPrecision_));
                            }
                        } else {
                            rawResults[row] = unscaled;
                        }
                    });
                } else {
                    rows.applyToSelected([&](vector_size_t row) {
                        rawResults[row] = static_cast<int128_t>(rawValues[row]);
                    });
                }
            }

            std::exception_ptr overflowException(
                int64_t unscaled,
                uint8_t resultPrecision) const {
                return std::make_exception_ptr(std::overflow_error(fmt::format(
                    "Unscaled value {} too large for precision {}.",
                    unscaled,
                    resultPrecision)));
            }

            const uint8_t resultPrecision_;
            const bool nullOnOverflow_;
        };

        std::shared_ptr<exec::VectorFunction> createMakeDecimal(
            const TypePtr &type,
            bool nullOnOverflow) {
            const uint8_t precision = getDecimalPrecisionScale(*type).first;
            if (type->isShortDecimal()) {
                return std::make_shared<MakeDecimalFunction<int64_t> >(
                    precision, nullOnOverflow);
            }
            return std::make_shared<MakeDecimalFunction<int128_t> >(
                precision, nullOnOverflow);
        }
    } // namespace

    TypePtr MakeDecimalCallToSpecialForm::resolveType(
        const std::vector<TypePtr> & /*argTypes*/) {
        POLLUX_FAIL("Make decimal function does not support type resolution.");
    }

    exec::ExprPtr MakeDecimalCallToSpecialForm::constructSpecialForm(
        const TypePtr &type,
        std::vector<exec::ExprPtr> &&args,
        bool trackCpuUsage,
        const core::QueryConfig & /*config*/) {
        POLLUX_USER_CHECK(
            type->isDecimal(),
            "The result type of make_decimal should be decimal type.");
        POLLUX_USER_CHECK_GE(
            args.size(), 1, "Make_decimal expects one or two arguments.");
        POLLUX_USER_CHECK_LE(
            args.size(), 2, "Make_decimal expects one or two arguments.");
        POLLUX_USER_CHECK_EQ(
            args[0]->type()->kind(),
            TypeKind::BIGINT,
            "The first argument of make_decimal should be of bigint type.");
        bool nullOnOverflow = true;
        if (args.size() > 1) {
            POLLUX_USER_CHECK_EQ(
                args[1]->type()->kind(),
                TypeKind::BOOLEAN,
                "The second argument of make_decimal should be of boolean type.");
            auto constantExpr = std::dynamic_pointer_cast<exec::ConstantExpr>(args[1]);
            POLLUX_USER_CHECK_NOT_NULL(
                constantExpr,
                "The second argument of make_decimal should be constant expression.");
            POLLUX_USER_CHECK(
                constantExpr->value()->is_constant_encoding(),
                "The second argument of make_decimal should be wrapped in constant vector.");
            auto constant_vector =
                    constantExpr->value()->as_unchecked<ConstantVector<bool> >();
            POLLUX_USER_CHECK(
                !constant_vector->is_null_at(0),
                "The second argument of make_decimal is non-nullable.");
            nullOnOverflow = constant_vector->value_at(0);
        }
        return std::make_shared<exec::Expr>(
            type,
            std::move(args),
            createMakeDecimal(type, nullOnOverflow),
            exec::VectorFunctionMetadata{},
            kMakeDecimal,
            trackCpuUsage);
    }
} // namespace kumo::pollux::functions::sparksql
